#pragma once

#include "updater_common.h"

namespace mdtx::online::updater
{

    // Linear decay
    struct lindecay_updater
    {
        double a;
        double decay;

        explicit lindecay_updater(double a)
            : a(a),
              decay(0.0)
        {
        }

        void update(double x)
        {
            decay = std::max(std::max(x, decay - a), 0.0);
        }
    };

    // Exponential decay
    struct expdecay_updater
    {
        double a;
        double decay;

        explicit expdecay_updater(double a)
            : a(a),
              decay(0.0)
        {
        }

        void update(double x)
        {
            decay = std::max(x, decay * a);
        }
    };

    template <class T>
    struct minmax_updater
    {
        boost::circular_buffer<T> mindeq, maxdeq;

        explicit minmax_updater(int max_window_size)
            : mindeq(max_window_size),
              maxdeq(max_window_size)
        {
        }

        void insert(T x)
        {
            while (!mindeq.empty() && mindeq.back() > x)
                mindeq.pop_back();
            while (!maxdeq.empty() && maxdeq.back() < x)
                maxdeq.pop_back();
            mindeq.push_back(x);
            maxdeq.push_back(x);
        }

        void remove(T x)
        {
            if (mindeq.front() == x)
                mindeq.pop_front();
            if (maxdeq.front() == x)
                maxdeq.pop_front();
        }

        void roll(T x, T y)
        {
            remove(y);
            insert(x);
        }
    };

    template <class T>
    struct argminmax_updater
    {
        boost::circular_buffer<T> mindeq, maxdeq;
        boost::circular_buffer<int> minidx, maxidx;

        explicit argminmax_updater(int max_window_size)
            : mindeq(max_window_size),
              maxdeq(max_window_size),
              minidx(max_window_size),
              maxidx(max_window_size)
        {
        }

        void insert(T x)
        {
            while (!mindeq.empty() && mindeq.back() > x)
            {
                mindeq.pop_back();
                minidx.pop_back();
            }
            while (!maxdeq.empty() && maxdeq.back() < x)
            {
                maxdeq.pop_back();
                minidx.pop_back();
            }
            for (auto &it : minidx)
                it += 1;
            for (auto &it : maxidx)
                it += 1;
            mindeq.push_back(x);
            maxdeq.push_back(x);
            minidx.push_back(0);
            maxidx.push_back(0);
        }

        void remove(T x)
        {
            if (mindeq.front() == x)
            {
                mindeq.pop_front();
                minidx.pop_front();
            }
            if (maxdeq.front() == x)
            {
                maxdeq.pop_front();
                maxidx.pop_front();
            }
        }

        void roll(T x, T y)
        {
            remove(y);
            insert(x);
        }
    };

    // Karnin, Z., Lang, K., & Liberty, E. (2016, October). Optimal quantile approximation in streams. In 2016 ieee 57th annual symposium on foundations of computer science (focs) (pp. 71-78). IEEE.
    // KLL sketch updater
    template <class T>
    class kll_updater
    {
        size_t _k, _size, _max_size;
        double _c;
        bool _lz;

        std::random_device _rdev;
        std::default_random_engine _reng;
        std::uniform_int_distribution<size_t> _runi;

        std::vector<std::vector<T>> compact;

        size_t level_cap(size_t lv) const
        {
            size_t d = compact.size() - lv - 1;
            return static_cast<size_t>(std::ceil(
                       std::pow(_c, d) * _k)) +
                   1;
        }

        void update_size()
        {
            _size = 0;
            for (const auto &cpt : compact)
            {
                _size += cpt.size();
            }
        }

        void update_max_size()
        {
            _max_size = 0;
            for (size_t lv = 0; lv < compact.size(); ++lv)
            {
                _max_size += level_cap(lv);
            }
        }

        void grow()
        {
            compact.push_back({});
            update_max_size();
        }

        void compress()
        {
            size_t maxlv = compact.size();
            for (size_t lv = 0; lv < maxlv; ++lv)
            {
                if (compact[lv].size() >= level_cap(lv))
                {
                    if (lv + 1 == maxlv)
                    {
                        grow();
                    }
                    // compress to next level
                    std::sort(compact[lv].begin(), compact[lv].end());
                    size_t offset = _runi(_reng);
                    for (size_t i = offset; i < compact[lv].size(); i += 2)
                    {
                        compact[lv + 1].push_back(compact[lv][i]);
                    }
                    compact[lv].clear();
                    update_size();
                    // only compress this level if lazy
                    if (_lz)
                    {
                        break;
                    }
                }
            }
        }

    public:
        kll_updater(size_t k, double c = 2.0 / 3.0, bool lazy = true)
            : _k(k),
              _size(0),
              _max_size(0),
              _c(c),
              _lz(lazy),
              _rdev(),
              _reng(_rdev()),
              _runi(0, 1),
              compact({})
        {
        }

        void insert(T x)
        {
            compact[0].push_back(x);
            if (++_size >= _max_size)
            {
                compress();
            }
        }

        std::pair<std::vector<T>, std::vector<double>> density(bool cumulative = false) const
        {
            size_t lv = 0;
            size_t w = 0;
            std::vector<std::pair<T, size_t>> weighted;

            weighted.reserve(_size);
            for (const auto &cpt_lv : compact)
            {
                const size_t w_lv = 1 << lv;
                for (auto val_lv : cpt_lv)
                {
                    weighted.emplace_back(val_lv, w_lv);
                    w += w_lv;
                }
                ++lv;
            }
            std::sort(weighted.begin(), weighted.end());

            std::vector<T> vals;
            std::vector<double> dens;
            for (const auto &item : weighted)
            {
                vals.push_back(item.first);
                dens.push_back(static_cast<double>(item.second) / w);
            }
            if (cumulative)
            {
                std::partial_sum(dens.begin(), dens.end(), dens.begin());
            }

            return std::make_pair(std::move(vals), std::move(dens));
        }

        std::vector<T> quantile(const std::vector<double> &probs) const
        {
            std::vector<T> q;

            auto dens = density(true);
            auto idx = dens.second.begin();
            auto idx_end = dens.second.end();
            for (auto p : probs)
            {
                idx = std::lower_bound(idx, idx_end, p);
                if (idx == idx_end)
                {
                    q.push_back(dens.first.back());
                }
                else
                {
                    q.push_back(dens.first[idx - dens.second.begin()]);
                }
            }

            return q;
        }

        void merge(const kll_updater &rhs)
        {
            while (compact.size() < rhs.compact.size())
            {
                grow();
            }
            // simply concatenate each compactor level
            for (size_t lv = 0; lv < rhs.compact.size(); ++lv)
            {
                compact[lv].insert(compact[lv].end(), rhs.compact[lv].begin(), rhs.compact[lv].end());
            }
            // compress if needed
            update_size();
            if (_size >= _max_size)
            {
                compress();
            }
        }
    };

    // Indexable skiplist
    constexpr size_t iskiplist_max_level = 32;
    template <class T, class CP = std::less<T>, class EQ = std::equal_to<T>>
    class iskiplist_updater
    {
    private:
        struct node_type;

        struct level_type
        {
            node_type *next;
            size_t span;
        };

        struct node_type
        {
            T v;
            std::vector<level_type> level;
            node_type(T value, size_t depth)
                : level(depth, {nullptr, 0}),
                  v(value)
            {
            }
        };

        node_type *mk_node(T value, size_t depth)
        {
            return new node_type(value, depth);
        }

        void rm_node(node_type *node)
        {
            delete node;
        }

        using fp_type = std::array<node_type *, iskiplist_max_level>;
        using fs_type = std::array<size_t, iskiplist_max_level>;
        void rm_node_fp(const fp_type &fp)
        {
            auto node = fp[0]->level[0].next;
            auto depth = node->level.size();
            for (size_t lv = 0; lv < depth; ++lv)
            {
                fp[lv]->level[lv].span += node->level[lv].span - 1;
                fp[lv]->level[lv].next = node->level[lv].next;
            }
            for (size_t lv = depth; lv < _mxlv; ++lv)
            {
                fp[lv]->level[lv].span -= 1;
            }
            rm_node(node);
            _size -= 1;
        }

        // Resource
        size_t _size, _mxlv;
        node_type *_head;

        // Random engine
        std::default_random_engine _rgen;
        std::uniform_int_distribution<> _runi;

        size_t random_depth()
        {
            size_t depth = 1;
            while (_runi(_rgen))
            {
                ++depth;
            }
            return std::min(depth, _mxlv);
        }

    public:
        explicit iskiplist_updater(size_t target_size)
            : _size(0),
              _mxlv(std::min(iskiplist_max_level, static_cast<size_t>(std::log2(target_size)) + 1)),
              _head(mk_node(T(), _mxlv)),
              _rgen((std::random_device())()),
              _runi(0, 1)
        {
        }

        ~iskiplist_updater()
        {
            node_type *node;
            while (_head)
            {
                node = _head;
                _head = _head->level[0].next;
                rm_node(node);
            }
        }

        // RM

        friend void swap(iskiplist_updater &lhs, iskiplist_updater &rhs) noexcept
        {
            std::swap(lhs._size, rhs._size);
            std::swap(lhs._mxlv, rhs._mxlv);
            std::swap(lhs._head, rhs._head);
        }

        iskiplist_updater(const iskiplist_updater &other)
            : iskiplist_updater(1)
        {
            merge(other);
        }

        iskiplist_updater(iskiplist_updater &&other) noexcept
            : _size(other._size),
              _mxlv(other._mxlv),
              _head(other._head),
              _rgen((std::random_device())()),
              _runi(0, 1)
        {
            // steal from other
            other._head = nullptr;
        }

        iskiplist_updater &operator=(iskiplist_updater rhs) noexcept
        {
            swap(*this, rhs);
            return *this;
        }

        // API

        size_t insert(T value)
        {
            fp_type front{};
            fs_type fspan{};
            // search frontier
            auto node = _head;
            for (size_t lv = _mxlv; lv--;)
            {
                while (node->level[lv].next &&
                       CP()(node->level[lv].next->v, value))
                {
                    fspan[lv] += node->level[lv].span;
                    node = node->level[lv].next;
                }
                front[lv] = node;
            }
            // create node
            auto depth = random_depth();
            auto newnode = mk_node(value, depth);
            // create link
            size_t bspan = 0;
            for (size_t lv = 0; lv < depth; ++lv)
            {
                auto prev = front[lv];
                newnode->level[lv].next = prev->level[lv].next;
                prev->level[lv].next = newnode;
                newnode->level[lv].span = prev->level[lv].span - bspan;
                prev->level[lv].span = bspan + 1;
                bspan += fspan[lv];
            }
            for (size_t lv = depth; lv < _mxlv; ++lv)
            {
                front[lv]->level[lv].span += 1;
                bspan += fspan[lv];
            }
            _size += 1;
            return bspan;
        }

        void remove(T value)
        {
            fp_type front{};
            size_t depth;
            auto node = _head;
            for (size_t lv = _mxlv; lv--;)
            {
                while (node->level[lv].next &&
                       CP()(node->level[lv].next->v, value))
                {
                    node = node->level[lv].next;
                }
                front[lv] = node;
            }
            if (node->level[0].next &&
                EQ()(node->level[0].next->v, value))
            {
                rm_node_fp(front);
            }
        }

        void remove_rank(size_t rank)
        {
            fp_type front{};
            size_t depth;
            auto node = _head;
            for (size_t lv = _mxlv; lv--;)
            {
                while (node->level[lv].next &&
                       rank >= node->level[lv].span)
                {
                    rank -= node->level[lv].span;
                    node = node->level[lv].next;
                }
                front[lv] = node;
            }
            rm_node_fp(front);
        }

        T operator[](size_t i) const
        {
            auto node = _head;
            i += 1; // ignore head
            for (size_t lv = _mxlv; lv--;)
            {
                while (node->level[lv].next &&
                       i >= node->level[lv].span)
                {
                    i -= node->level[lv].span;
                    node = node->level[lv].next;
                }
            }
            return node->v;
        }

        T at(size_t i) const
        {
            if (i > _size)
            {
                throw std::out_of_range("Index out of range.");
            }
            return this->operator[](i);
        }

        size_t rank(T value) const
        {
            auto node = _head;
            size_t rank = 0;
            for (size_t lv = _mxlv; lv--;)
            {
                while (node->level[lv].next &&
                       CP()(node->level[lv].next->v, value))
                {
                    rank += node->level[lv].span;
                    node = node->level[lv].next;
                }
            }
            return rank;
        }

        void merge(const iskiplist_updater &other)
        {
            size_t tgt_size = std::max(1U << this->_mxlv, 1U << other._mxlv) - 1;
            size_t new_size = std::max(this->_size + other._size, tgt_size);
            iskiplist_updater new_list(new_size);

            fp_type front;
            std::fill(front.begin(), front.end(), new_list._head);

            auto nodeL = this->_head->level[0].next;
            auto nodeR = other._head->level[0].next;
            while (nodeL || nodeR)
            {
                if (!nodeL)
                {
                    std::swap(nodeL, nodeR);
                }
                if (nodeR && CP()(nodeR->v, nodeL->v))
                {
                    std::swap(nodeL, nodeR);
                }
                auto depth = new_list.random_depth();
                auto node = new_list.mk_node(nodeL->v, depth);
                for (size_t lv = 0; lv < new_list._mxlv; ++lv)
                {
                    front[lv]->level[lv].span += 1;
                }
                for (size_t lv = 0; lv < depth; ++lv)
                {
                    front[lv]->level[lv].next = node;
                    front[lv] = node;
                }
                new_list._size += 1;
                nodeL = nodeL->level[0].next;
            }
            std::swap(*this, new_list);
        }

        size_t size() const
        {
            return _size;
        }

        bool empty() const
        {
            return _size == 0;
        }

        // ITERATOR
        class iterator
        {
        private:
            node_type *node;

        public:
            friend class iskiplist_updater;

            using value_type = T;
            using pointer = T const *;
            using reference = const T &;
            using difference_type = std::ptrdiff_t;
            using iterator_category = std::input_iterator_tag;

            reference operator*() const
            {
                return node->v;
            }

            pointer operator->() const
            {
                return &(node->v);
            }

            iterator &operator++()
            {
                node = node->level[0].next;
                return *this;
            }

            iterator operator++(int)
            {
                iterator it(*this);
                node = node->level[0].next;
                return it;
            }

            bool operator==(const iterator &rhs) const
            {
                return node == rhs.node;
            }

            bool operator!=(const iterator &rhs) const
            {
                return node != rhs.node;
            }
        };

        iterator begin()
        {
            iterator it;
            it.node = _head->level[0].next;
            return it;
        }

        iterator end()
        {
            iterator it;
            it.node = nullptr;
            return it;
        }

        iterator find(T value)
        {

            auto node = _head;
            for (size_t lv = _mxlv; lv--;)
            {
                while (node->level[lv].next &&
                       CP()(node->level[lv].next->v, value))
                {
                    node = node->level[lv].next;
                }
            }

            iterator it;
            if (node->level[0].next &&
                EQ()(node->level[0].next->v, value))
            {
                it.node = node->level[0].next;
            }
            else
            {
                it.node = nullptr;
            }

            return it;
        }
    };

#ifdef MDTX_USE_ARMADILLO

    struct rls_updater
    {
        arma::colvec w;
        arma::mat P;
        double lamb, ilamb;

        rls_updater(int size, double lambda = 0.98, double sigma = 1.0)
            : w(size, arma::fill::zeros),
              P(size, size, arma::fill::zeros),
              lamb(lambda),
              ilamb(1.0 / lambda)
        {
            P.diag().fill(sigma);
        }

        void update(const arma::rowvec &x, double d)
        {
            double a = d - arma::dot(x, w);
            arma::mat Px{P * x.t()};
            arma::colvec g{Px / (lamb + arma::dot(x, Px))};
            P *= ilamb * (1 - arma::dot(g, x));
            w += a * g;
        }

        double predict(const arma::rowvec &x)
        {
            return arma::dot(x, w);
        }
    };

    struct nlms_updater
    {
        double _mu, _eps;
        arma::colvec w;

        nlms_updater(int size, double mu, double eps = 1e-11)
            : _mu(mu),
              _eps(eps),
              w(size, arma::fill::zeros)
        {
        }

        void update(const arma::rowvec &x, double d)
        {
            double a = d - arma::dot(x, w);
            w += _mu * a * x.t() / (arma::dot(x, x) + _eps);
        }

        double predict(const arma::rowvec &x)
        {
            return arma::dot(x, w);
        }
    };

    struct kmeans_updater
    {
        int k, n;
        double hl;
        arma::mat cent;
        arma::colvec dist;

        kmeans_updater(int ncluster, int ndim, double halflife)
            : k(ncluster),
              n(0),
              hl(halflife),
              cent(k, ndim, arma::fill::zeros),
              dist(k, arma::fill::zeros)
        {
        }

        int update(const arma::rowvec &x)
        {
            int idx = n;
            if (n < k)
            {
                arma::rowvec tmp(x.n_cols, arma::fill::randu);
                cent.row(n) = x % tmp;
                n += 1;
            }
            else
            {
                idx = predict(x);
                cent.row(idx) += hl * (x - cent.row(idx));
            }
            return idx;
        }

        int predict(const arma::rowvec &x)
        {
            for (int i = 0; i < n; ++i)
            {
                dist(i) = arma::norm(cent.row(i) - x);
            }
            return dist.index_min();
        }
    };

#endif

}
